# -*- coding: utf-8 -*-
"""
@Time    : 2024/8/21 18:54 
@Author  : ZhangShenao 
@File    : 4.使用StructuredTool类方法自定义工具.py 
@Desc    : 使用StructuredTool类方法自定义工具

针对一个已有的函数,通过StructuredTool.from_function()类方法,可以快速创建自定义工具
该方式支持同时创建同步和异步的工具
可以为工具指定name、description等属性

对于已有的函数,并且想同时创建同步和异步工具,就可以考虑使用StructuredTool.from_function()这种方式
"""
from langchain_core.tools import StructuredTool

from calculator_input import CalculatorInput


# 创建原始函数
def multiply(x: int, y: int) -> int:
    """
    同步方法,计算两个整数的乘积
    :param x: 第一个整数
    :param y: 第二个整数
    :return: 两个数的乘积
    """
    return x * y


async def amultiply(x: int, y: int) -> int:
    """
    异步方法,计算两个整数的乘积
    :param x: 第一个整数
    :param y: 第二个整数
    :return: 两个数的乘积
    """
    return x * y


# 使用StructuredTool.from_function()类方法,创建自定义工具
multiply_tool = StructuredTool.from_function(
    func=multiply,  # 指定同步工具的函数
    coroutine=amultiply,  # 指定异步工具的函数
    name="multiply_tool",  # 指定工具的名称
    description="计算两个整数的乘积",  # 指定工具的描述信息
    args_schema=CalculatorInput,  # 指定工具的参数结构
    return_direct=True,  # 设置工具为直接返回结果
)

# 获取工具的属性信息
print(f"工具名称: {multiply_tool.name}")
print(f"工具描述: {multiply_tool.description}")
print(f"工具参数定义: {multiply_tool.args}")
print(f"工具是否直接返回: {multiply_tool.return_direct}")

# 调用工具,获取结果
result = multiply_tool.invoke({"x": 3, "y": 4})
print(f"工具调用结果: {result}")
